# coding:UTF-8
import numpy as np
from ensembel.decision_tree import predict
import pickle
from ensembel.random_forest import random_forest_training

def load_data(file_name):
    '''
    导入数据
    :param file_name:训练数据保存的文件名
    :return: 训练数据
    '''
    data_train = []
    f = open(file_name)
    for line in f.readlines():
        lines = line.strip().split("\t")
        data_tmp = []
        for x in lines:
            data_tmp.append(float(x))
        data_train.append(data_tmp)
    f.close()
    return data_train

def get_predict(trees_result, trees_feature, data_train):
    '''
    利用训练好的随机森林模型对样本进行预测
    :param trees_result: 训练好的随机森林模型
    :param trees_feature: 每一棵树选择的特征
    :param data_train: 训练样本
    :return: 对样本预测的结果
    '''
    m_tree = len(trees_result)
    m = np.shape(data_train)[0]

    result = []
    for i in range(m_tree):
        clf = trees_result[i]
        feature = trees_feature[i]
        data = split_data(data_train, feature)
        result_i = []
        for i in range(m):
            result_i.append(list(predict(data[i][0:-1], clf).keys())[0])
        result.append(result_i)
    final_predict = np.sum(result, axis=0)
    return final_predict

def cal_correct_rate(data_train, final_predict):
    '''
    计算机模型的预测准确性
    :param data_train: 训练样本
    :param final_predict: 预测结果
    :return: 准确性
    '''
    m = len(final_predict)
    corr = 0.0
    for i in range(m):
        if data_train[i][-1] * final_predict[i] > 0:
            corr += 1
    return corr / m

def save_model(trees_result, trees_feature, result_file, feature_file):
    '''
    保存最终的模型
    :param trees_result:训练好的随机森林模型
    :param trees_feature: 每一棵决策树选择的特征
    :param result_file: 模型保存的文件
    :param feature_file: 特征被保存的文件
    '''

    m = len(trees_feature)
    f_fea = open(feature_file, "w")
    for i in range(m):
        fea_tmp = []
        for x in trees_feature[i]:
            fea_tmp.append(str(x))
        f_fea.writelines("\t".join(fea_tmp) + "\n")
    f_fea.close()

    with open(result_file, "wb") as f:
        pickle.dump(trees_result, f)

def split_data(data_train, feature):
    '''
    选择特征
    :param data_train:训练数据集
    :param feature: 要选择的数据集
    :return: 选择出来的数据集
    '''
    m = np.shape(data_train)[0]
    data = []

    for i in range(m):
        data_x_tmp = []
        for x in feature:
            data_x_tmp.append(data_train[i][x])
        data_x_tmp.append(data_train[i][-1])
        data.append(data_x_tmp)
    return data

if __name__ == "__main__":
    print("------1. load data")
    pwd = "/home/xiefeihong/PycharmProjects/SimpleMachineLearning/static/ensembel/"
    data_train = load_data(pwd + "data.txt")

    print("------2. random forest training")
    trees_result, trees_feature = random_forest_training(data_train, 50)

    print("------3. get prediction correct rate")
    result = get_predict(trees_result, trees_feature, data_train)
    corr_rate = cal_correct_rate(data_train, result)
    print("\t------correct rate: ", corr_rate)

    print("------4. save model")
    save_model(trees_result, trees_feature, pwd + "result_file", pwd + "feature_file")